import os

os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"


import json
import logging
import time
from collections import namedtuple
from functools import partial

import jax.numpy as jnp
import numpy as np
import optax
import orbax.checkpoint
from flax import nnx
from flax.training.train_state import TrainState

from model.diffusion.diffusion import DiffusionPolicy
from model.networks.mlp import DiffusionMLP
from model.utils.jax_types import Params

Batch = namedtuple("Batch", "rotations positions")


# class TwoLayerMLP(nnx.Module):
#   def __init__(self, rngs: nnx.Rngs):
#     self.linear1 = nnx.Linear(6, 256, rngs=rngs, use_bias=True)
#     self.linear2 = nnx.Linear(256, 256, rngs=rngs, use_bias=True)
#     self.linear3 = nnx.Linear(256, 6, rngs=rngs, use_bias=True)

#   def __call__(self, x):
#     x = self.linear1(x)
#     x = jax.nn.mish(x)
#     x = self.linear2(x)
#     x = jax.nn.mish(x)
#     return self.linear3(x)

# # Instantiate the model and show we can run it.
# model = TwoLayerMLP(rngs=nnx.Rngs(0))
# x = jax.random.normal(jax.random.key(42), (256, 6))


# for i in range(10000):
#     model(x)
ROTATION_MAX = np.array([180.0000, 89.1460, 180.0000])
ROTATION_MIN = np.array([-179.9690, -8.3270, -179.9670])

JOINT_MAX = np.array([150, 180, 0, 100, 70, 120])
JOINT_MIN = np.array([-150, 0, -170, -100, -70, -120])


def denormalize_gripper_rot(r_xyz):
    return (np.array(r_xyz) + 1) / 2 * (ROTATION_MAX - ROTATION_MIN) + ROTATION_MIN


class DiffusionGripperModel(DiffusionPolicy):
    """
    Model:
    map gripper locations (x, y, z) to (j1, j2, j3, j4, j5, j6)
    """

    def __init__(self, seed=0, model_file: str = None, predict_epsilon: bool = False):

        action_dim = 3  # (rx, ry, rz)
        horizon_steps = 1
        obs_dim = 3  # condition on (x, y, z)

        network = DiffusionMLP(output_dim=action_dim, activation_type="Mish", mlp_dims=(256, 256, 256))

        super().__init__(
            seed=seed,
            model=network,
            action_dim=action_dim,
            horizon_steps=horizon_steps,
            obs_dim=obs_dim,
            max_grad_norm=1,
            lr=3e-4,
            predict_epsilon=predict_epsilon,
        )

        if model_file is not None:
            self.load_ckpt(model_file)

    def __call__(self, gripper_pos):
        if gripper_pos.ndim == 1:
            # add batch dim
            gripper_pos = gripper_pos[None, :]

        gripper_rot = self.sampling(obs=gripper_pos).reshape(len(gripper_pos), -1)
        return jnp.concat([gripper_pos, denormalize_gripper_rot(gripper_rot)], axis=-1)

    def train(self, dataloader, num_epochs: int, save_as: str):
        from tensorboardX import SummaryWriter

        writer = SummaryWriter(log_dir=f"log/gj_map/{round(time.time())}")

        avg_loss = 0

        for epoch in range(num_epochs):

            cnt_batch = 0
            loss_rec = []

            for batch in dataloader:

                # info = self.update(batch.joint, batch.gripper)

                info = self.update(batch.rotations, batch.positions)

                loss_rec.append(info["loss"].item())
                cnt_batch += 1

            avg_loss = sum(loss_rec) / len(loss_rec)
            writer.add_scalar("train/loss", avg_loss, epoch)

            if (1 + epoch) % 100 == 0:
                print(f"Epoch{epoch+1}: loss = {avg_loss}")
                self.save_ckpt(file_name=save_as)

        self.save_ckpt(file_name=save_as)

    def save_ckpt(self, save_dir: str = "data/weights", file_name: str = "nn_3d_6d_diffusion.ckp"):

        if not os.path.exists(save_dir):
            os.mkdir(save_dir)

        filename = os.path.join(os.getcwd(), save_dir, file_name)

        orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
        orbax_checkpointer.save(filename, {"params": self.network.params}, force=True)
        print(f"Successfully save model to {filename}")

    def load_ckpt(self, file_path: str = "data/weights/gp_mapping.ckp"):
        orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
        filename = os.path.join(os.getcwd(), file_path)
        raw_restored = orbax_checkpointer.restore(filename)
        self.network = self.network.replace(params=raw_restored["params"])

        # run it one time to save the structure in the memory
        # self(jnp.empty((3,)))

        print(f"Successfully load model from {filename}")


if __name__ == "__main__":
    import torch
    from torch.utils.data import Dataset

    from model.data_processing.jax_dataloader import NumpyLoader

    class GripperData(Dataset):
        def __init__(self, filename):

            with open(filename, "r") as f:
                data = json.load(f)

                self.gripper_pos = torch.tensor(data["gripper"])
                self.gripper_xyz = self.gripper_pos[:, :3]
                self.gripper_rot = self.gripper_pos[:, 3:]

                # normalize the target to [-1, 1]
                self.gripper_rot = (self.gripper_rot - ROTATION_MIN) * 2 / (ROTATION_MAX - ROTATION_MIN) - 1

        def __getitem__(self, idx):
            """
            condition: gripper pos
            action: joint values
            """
            return Batch(self.gripper_rot[idx], self.gripper_xyz[idx])

        def __len__(self):
            return len(self.gripper_rot)

    data = GripperData("data/motions/full_gj.json")
    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )

    self = DiffusionGripperModel(predict_epsilon=False)
    # self.load_ckpt("data/weights/nn_3d_6d_full_space_v1.ckp")
    self.train(data_loader, 10000, "nn_3d_6d_diffusion_rot_val.ckp")
